`timescale 1ns / 1ps

`define SLICE(BIT, IDX) (``BIT``) * (``IDX``) + (``BIT``) - 1 : (``BIT``) * (``IDX``)

module maxpool2x2
#(
    parameter N_IO      = 4,
    parameter N_CH      = 48,
    parameter N_COL     = 320,
    parameter BIT       = 8,
    parameter FIFOA_TYPE = "distributed",
    parameter FIFOA_LATENCY = 1,
    parameter FIFOB_TYPE = "block",
    parameter FIFOB_LATENCY = 2
)
(
    input   clk,
    input   rst,
    
    input   i_vld,
    output  i_rdy,
    input   [N_IO*BIT-1 : 0]      i_data,
    
    output  o_vld,
    input   o_rdy,
    output  [N_IO*BIT-1 : 0]      o_data
);

localparam WIDTH = N_IO * BIT;
localparam FOLD = N_CH / N_IO;

reg  row_lsb;

wire [1: 0] cstate;
reg         enable;
reg         ii_rdy;

wire [$clog2(N_COL)-1 : 0]  col;
wire                        col_ena;
wire                        col_last;
wire                        fod_last;

wire                ai_rdy;
reg                 ai_vld;
wire [WIDTH-1 : 0]  ai_data;
reg                 ao_rdy;
wire                ao_vld;
wire [WIDTH-1 : 0]  ao_data;

wire                bi_rdy;
reg                 bi_vld;
wire [WIDTH-1 : 0]  bi_data;
reg                 bo_rdy;
wire                bo_vld;
wire [WIDTH-1 : 0]  bo_data;

wire [WIDTH-1 : 0]  mx_ain;
wire [WIDTH-1 : 0]  mx_bin;
wire [WIDTH-1 : 0]  mx_out;

wire                r_rdy;
reg                 r_vld;
reg  [WIDTH-1 : 0]  r_data;

assign cstate = { row_lsb, col[0] };
assign col_ena = enable & fod_last;

assign ai_data = cstate[1] ? mx_out : i_data;
assign bi_data = mx_out;

assign mx_ain = i_data;
assign mx_bin = cstate[0] ? ao_data : bo_data;

assign r_rdy = ~r_vld | o_rdy;
assign i_rdy = ii_rdy;
assign o_vld  = r_vld;
assign o_data = r_data;

always @(*)
begin
    (* full_case *)
    case (cstate)
		2'b00:
        begin
            ii_rdy = r_rdy & ai_rdy;
            enable = r_rdy & ai_rdy & i_vld;
            ai_vld = enable;
            ao_rdy = 1'b0;
            bi_vld = 1'b0;
            bo_rdy = 1'b0;
		end
		2'b01:
        begin
            ii_rdy = r_rdy & ao_vld & bi_rdy;
            enable = r_rdy & ao_vld & bi_rdy & i_vld;
            ai_vld = 1'b0;
            ao_rdy = enable;
            bi_vld = enable;
            bo_rdy = 1'b0;
		end
		2'b10:
        begin
            ii_rdy = r_rdy & ai_rdy & bo_vld;
            enable = r_rdy & ai_rdy & bo_vld & i_vld;
            ai_vld = enable;
            ao_rdy = 1'b0;
            bi_vld = 1'b0;
            bo_rdy = enable;
		end
		2'b11:
        begin
            ii_rdy = r_rdy & ao_vld;
            enable = r_rdy & ao_vld & i_vld;
            ai_vld = 1'b0;
            ao_rdy = enable;
            bi_vld = 1'b0;
            bo_rdy = 1'b0;
		end
	endcase
end

always @(posedge clk)
begin
    if (rst) begin
          r_vld <= 1'b0;
        row_lsb <= 1'b0;
    end
    else begin
        row_lsb <= row_lsb ^ (enable & col_last & fod_last);
        if (r_rdy)
            r_vld <= (cstate == 2'b11) & ao_vld & i_vld;
        if (enable & cstate == 2'b11)
            r_data <= mx_out;
    end
end

zq_counter #(
    .N  (FOLD)
) inst_cnt_fold
(
    .clk    (clk),
    .rst    (rst),
    .clken  (enable),
    .last   (fod_last),
    .out    ()
);

zq_counter #(
    .N  (N_COL)
) inst_cnt_col
(
    .clk    (clk),
    .rst    (rst),
    .clken  (col_ena),
    .last   (col_last),
    .out    (col)
);

zq_max #(
    .N_IO   ( N_IO  ),
    .BIT    ( BIT   )
)
inst_zqmax (
    .ain    ( mx_ain ),
    .bin    ( mx_bin ),
    .out    ( mx_out )
);

zq_fifo #(
    .WIDTH      ( WIDTH         ),
    .DEPTH      ( FOLD          ),
    .RAMTYPE    ( FIFOA_TYPE    ),
    .RAMLATENCY ( FIFOA_LATENCY )
)
inst_fifoA (
    .clk                     ( clk      ),
    .rst                     ( rst      ),

    .in_vld                  ( ai_vld   ),
    .in_rdy                  ( ai_rdy   ),
    .din                     ( ai_data  ),

    .out_vld                 ( ao_vld   ),
    .out_rdy                 ( ao_rdy   ),
    .dout                    ( ao_data  )
);

zq_fifo #(
    .WIDTH      ( WIDTH             ),
    .DEPTH      ( FOLD * N_COL / 2  ),
    .RAMTYPE    ( FIFOB_TYPE        ),
    .RAMLATENCY ( FIFOB_LATENCY     )
)
inst_fifoB (
    .clk                     ( clk      ),
    .rst                     ( rst      ),

    .in_vld                  ( bi_vld   ),
    .in_rdy                  ( bi_rdy   ),
    .din                     ( bi_data  ),

    .out_vld                 ( bo_vld   ),
    .out_rdy                 ( bo_rdy   ),
    .dout                    ( bo_data  )
);

endmodule


module zq_max
#(
    parameter N_IO      = 4,
    parameter BIT       = 8
)
(
    input   [N_IO*BIT-1 : 0]      ain,
    input   [N_IO*BIT-1 : 0]      bin,
    
    output  [N_IO*BIT-1 : 0]      out
);

wire [BIT-1 : 0] x [N_IO-1 : 0];
wire [BIT-1 : 0] y [N_IO-1 : 0];
wire [BIT-1 : 0] z [N_IO-1 : 0];

genvar i;
generate
    for (i = 0; i < N_IO; i = i + 1)
    begin
        assign x[i] = ain[ `SLICE(BIT, i) ];
        assign y[i] = bin[ `SLICE(BIT, i) ];
        assign z[i] = x[i] > y[i] ? x[i] : y[i];
        assign out[ `SLICE(BIT, i) ] = z[i];
    end
endgenerate

endmodule
